package cn.iocoder.yudao.framework.websocket.core.interceptor;

import cn.hutool.core.util.ObjectUtil;
import cn.iocoder.yudao.framework.security.config.SecurityProperties;
import cn.iocoder.yudao.framework.security.core.LoginUser;
import cn.iocoder.yudao.framework.security.core.util.SecurityFrameworkUtils;
import cn.iocoder.yudao.framework.tenant.core.context.TenantContextHolder;
import cn.iocoder.yudao.framework.websocket.config.WebSocketProperties;
import cn.iocoder.yudao.framework.websocket.core.WebSocketUser;
import cn.iocoder.yudao.module.system.api.oauth2.OAuth2TokenApi;
import cn.iocoder.yudao.module.system.api.oauth2.dto.OAuth2AccessTokenCheckRespDTO;
import lombok.extern.slf4j.Slf4j;
import org.springframework.http.server.ServerHttpRequest;
import org.springframework.http.server.ServletServerHttpRequest;
import org.springframework.messaging.Message;
import org.springframework.messaging.MessageChannel;
import org.springframework.messaging.simp.stomp.StompCommand;
import org.springframework.messaging.simp.stomp.StompHeaderAccessor;
import org.springframework.messaging.support.ChannelInterceptor;
import org.springframework.messaging.support.MessageHeaderAccessor;
import org.springframework.util.Assert;
import org.springframework.web.socket.handler.AbstractWebSocketHandler;

import javax.servlet.http.HttpServletRequest;
import java.security.Principal;
import java.util.List;
import java.util.Map;
import java.util.Objects;

@Slf4j
public class WebSocketChannelInterceptor implements ChannelInterceptor {
    private final OAuth2TokenApi oauth2TokenApi;
    private final SecurityProperties securityProperties;

    private final WebSocketProperties webSocketProperties;
    public WebSocketChannelInterceptor(OAuth2TokenApi oauth2TokenApi,
                                       SecurityProperties securityProperties,
                                       WebSocketProperties webSocketProperties){
        this.oauth2TokenApi = oauth2TokenApi;
        this.securityProperties = securityProperties;
        this.webSocketProperties = webSocketProperties;
    }
    /**
     * 发送消息前可以做的业务
     *
     */
    @Override
    public Message<?> preSend(Message<?> message, MessageChannel channel) {
        StompHeaderAccessor accessor = MessageHeaderAccessor.getAccessor(message, StompHeaderAccessor.class);
        //1、判断是否首次连接
        Assert.notNull(accessor,"websocket发生严重错误");
        Map<String, Object> sessionAttributes = accessor.getSessionAttributes();
        Assert.notNull(sessionAttributes,"sessionAttributes 不可能为空");
        ServerHttpRequest request = (ServerHttpRequest)sessionAttributes.get(WebSocketHandshakeInterceptor.HTTP_REQUEST);
        HttpServletRequest serverHttpRequest = ((ServletServerHttpRequest) request).getServletRequest();
        if (StompCommand.CONNECT.equals(accessor.getCommand())) {
            log.info("第一次连接");
            if(!webSocketProperties.isDebug()){
                List<String> nativeHeader = accessor.getNativeHeader(securityProperties.getTokenHeader());
                Assert.notNull(nativeHeader,"开启了websocket验证模式，websocket连接需要在header中携带当前用户的token");
                String customHeader = nativeHeader.get(0);
                OAuth2AccessTokenCheckRespDTO accessToken = oauth2TokenApi.checkAccessToken(customHeader).getCheckedData();
                LoginUser loginUser = new WebSocketUser(Objects.requireNonNull(accessor.getSessionId())).setId(accessToken.getUserId()).setUserType(accessToken.getUserType())
                        .setTenantId(accessToken.getTenantId()).setScopes(accessToken.getScopes());
                accessor.setUser(loginUser);
            }
        }
        else if(StompCommand.SUBSCRIBE.equals(accessor.getCommand())) {
            log.info(String.valueOf(accessor));
        }
        Principal user = accessor.getUser();
        if(!ObjectUtil.isNull(user)){
            WebSocketUser loginUser = (WebSocketUser)user;
            SecurityFrameworkUtils.setLoginUser(loginUser,serverHttpRequest);
            TenantContextHolder.setTenantId(loginUser.getTenantId());
        }
        LoginUser loginUser = SecurityFrameworkUtils.getLoginUser();
        log.info("{}",loginUser);
        //不是首次连接，已经登陆成功
        return message;
    }

}
